import os
import email
import pandas as pd
from tqdm import tqdm
from cot_utils import generate_cot

def read_email_content(file_path):
    """读取邮件内容，返回文本部分"""
    try:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            msg = email.message_from_file(f)
            
        # 获取邮件主体
        body = ""
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == "text/plain":
                    body += part.get_payload(decode=True).decode('utf-8', errors='ignore')
        else:
            body = msg.get_payload(decode=True).decode('utf-8', errors='ignore')
        
        # 清理和简化文本
        body = body.replace('\n', ' ').replace('\r', ' ').strip()
        return body[:1000]  # 限制长度
    except Exception as e:
        print(f"Error reading {file_path}: {str(e)}")
        return ""

def process_directory(directory, label):
    """处理指定目录下的所有邮件"""
    emails = []
    labels = []
    
    for filename in tqdm(os.listdir(directory), desc=f"Processing {label} emails"):
        if filename.startswith('.'):  # 跳过隐藏文件
            continue
        file_path = os.path.join(directory, filename)
        if os.path.isfile(file_path):
            content = read_email_content(file_path)
            if content:  # 只添加非空内容
                emails.append(content)
                labels.append(label)
    
    return emails, labels

# 处理ham邮件
print("处理正常邮件...")
ham_dir = "data/spamassassin/easy_ham"
ham_emails, ham_labels = process_directory(ham_dir, "ham")

# 处理spam邮件
print("处理垃圾邮件...")
spam_dir = "data/spamassassin/spam"
spam_emails, spam_labels = process_directory(spam_dir, "spam")

# 合并数据
emails = ham_emails + spam_emails
labels = ham_labels + spam_labels

# 创建DataFrame
df = pd.DataFrame({
    'text': emails,
    'label': labels
})

# 生成思维链提示
print("生成思维链提示...")
df['cot'] = df.apply(lambda row: generate_cot(row['text'], row['label']), axis=1)

# 显示基本信息
print("\n数据集信息:")
print(df['label'].value_counts())
print("\n示例:")
print(df.head())

# 保存为CSV
output_file = 'data/spamassassin_processed.csv'
df.to_csv(output_file, index=False)
print(f"\n数据已保存至: {output_file}") 